package com.smh;

import org.junit.jupiter.api.Test;

/**
 * @author shiminghui
 * @date 2025/3/1 10:20
 * @description: TODO
 */
public class _035_红黑树 {

    @Test
    public void test() {
        RBTree<Integer> rbTree = new RBTree<>();
        rbTree.put(4, 4);
        rbTree.put(2, 2);
        rbTree.put(8, 8);
        rbTree.put(7, 7);
        rbTree.put(9, 9);
        rbTree.put(5, 5);
        rbTree.put(6, 6);
        rbTree.remove(7);
    }


    private static class RBTree<T> {
        private RBNode<T> root;

        public RBTree() {
        }

        /**
         * 根据key删除节点
         *
         * @param key
         * @return
         */
        public T remove(int key) {
            RBNode<T> trbNode = find(key);
            if (trbNode == null) {
                return null;
            }
            // 找到节点了，开始删除操作
            doRemove(trbNode);
            return trbNode.value;
        }

        /**
         * 递归删除节点，当删除节点们没有孩子，和只有一个孩子的时候，直接删除，并调整树。当删除节点有两个孩子的时候，
         * 这时候要找到后继节点，然后交换删除节点和后继节点的键和值，然后递归删除后继节点。
         *
         * @param deleteNode
         */
        private void doRemove(RBNode<T> deleteNode) {
            RBNode<T> replaceNode = findReplaceNode(deleteNode);
            RBNode<T> parent = deleteNode.parent;

            if (replaceNode == null) { // 删除的节点是叶子节点，直接删除即可
                // case1：删除的节点是根节点
                if (deleteNode == root) { // 删除的节点是根节点，直接删除即可
                    root = null;
                    return;
                }
                // 调整树结构
                if (isBlack(deleteNode)) {
                    fixDoubleBlack(deleteNode);
                }

                if (deleteNode.isLeft()) {
                    parent.left = null;
                } else {
                    parent.right = null;
                }
                deleteNode.parent = null; //垃圾回收
                return;
            }
            // 当删除节点有一个孩子
            if (deleteNode.left == null || deleteNode.right == null) {
                // case1：删除的节点是根节点
                if (deleteNode == root) {
                    // 这里删除的是根节点，还只有一个孩子，那她的孩子肯定是红色，且他的孩子是叶子节点
                    deleteNode.key = replaceNode.key;
                    deleteNode.value = replaceNode.value;
                    deleteNode.left = deleteNode.right = null;
                    replaceNode.parent = null; // 垃圾回收
                    return;
                }

                // 删除节点是父亲节点的左孩子
                if (deleteNode.isLeft()) {
                    parent.left = replaceNode;
                }
                // 删除节点是父亲节点的右孩子
                else if (deleteNode.isRight()) {
                    parent.right = replaceNode;
                }
                replaceNode.parent = parent;
                deleteNode.parent = deleteNode.left = deleteNode.right = null; // 垃圾回收

                // 调整红黑树
                if (isBlack(deleteNode) && isBlack(replaceNode)) {
                    // 删除的节点是黑色，替换节点也是黑色，少了一个黑色
                    fixDoubleBlack(replaceNode);
                } else {
                    // case2：删除的节点是黑色，后继节点是红色
                    replaceNode.color = Color.RED;
                }

                return;
            }
            // case0：删除的节点有两个孩子
            // 让交换删除节点和后继节点的键和值，然后删除后继节点
            int key = replaceNode.key;
            replaceNode.key = deleteNode.key;
            deleteNode.key = key;

            T value = replaceNode.value;
            replaceNode.value = deleteNode.value;
            deleteNode.value = value;

            doRemove(replaceNode);// 删除后继节点
        }

        /**
         * 调整红黑树，双黑的问题，即删除的节点为黑色，剩余的节点也是黑色
         *
         * @param node
         */
        private void fixDoubleBlack(RBNode<T> node) {
            if (node == root) {
                return;
            }
            RBNode<T> parent = node.parent;
            RBNode<T> sibling = node.sibling();
            // case3：删除节点的兄弟是红色
            if (isRed(sibling)) {
                // 调整
                if (node.isLeft()) {
                    // 左旋
                    leftRotate(parent);
                } else if (node.isRight()) {
                    // 右旋
                    rightRotate(parent);
                }
                parent.color = Color.RED;
                sibling.color = Color.BLACK;

                fixDoubleBlack(node);
                return;
            }
            if (sibling != null) {
                // case4：删除节点的兄弟是黑色，并且侄子也是黑色
                if (isBlack(sibling.left) && isBlack(sibling.right)) {
                    sibling.color = Color.RED; // 兄弟变红
                    // 然后父亲变黑，使整个树平衡
                    if (isBlack(parent)) {
                        fixDoubleBlack(parent);
                    } else {
                        parent.color = Color.BLACK;
                    }
                }
                // case5：删除节点的兄弟是黑色，并且侄子是红色,至少一个
                else if (isBlack(sibling.left) && (isRed(sibling.right) || isRed(sibling.left))) {
                    //这里有四种情况
                    // case1：兄弟节点是左孩子，并且左侄子是红色 LL
                    if (sibling.isLeft() && isRed(sibling.left)) {
                        sibling.left.color = Color.BLACK;
                        sibling.color = parent.color;
                        parent.color = Color.BLACK;
                        rightRotate(parent);
                    }
                    // case2：兄弟节点是右孩子，并且右侄子是红色 RR
                    else if (sibling.isRight() && isRed(sibling.right)) {
                        sibling.right.color = Color.BLACK;
                        sibling.color = parent.color;
                        parent.color = Color.BLACK;
                        leftRotate(parent);
                    }
                    // case3：兄弟节点是左孩子，并且右侄子是红色 LR
                    else if (sibling.isLeft() && isRed(sibling.right)) {
                        sibling.color = Color.BLACK;
                        sibling.right.color = parent.color;
                        parent.color = Color.BLACK;
                        // 先左旋兄弟节点
                        leftRotate(sibling);
                        // 再右旋父亲节点
                        rightRotate(parent);
                    } else {
                        // case4：兄弟节点是右孩子，并且左侄子是红色 RL
                        sibling.color = Color.BLACK;
                        sibling.left.color = parent.color;
                        parent.color = Color.BLACK;
                        // 先右旋兄弟节点
                        rightRotate(sibling);
                        // 再左旋父亲节点
                        leftRotate(parent);
                    }
                }
            } else {
                fixDoubleBlack(parent);
            }
        }

        /**
         * 找到删除后剩余的节点
         *
         * @param node
         * @return
         */
        private RBNode<T> findReplaceNode(RBNode<T> node) {
            if (node.left == null && node.right == null) {
                return null;
            }
            if (node.left == null) {
                return node.right;
            }
            if (node.right == null) {
                return node.left;
            }

            RBNode<T> cur = node.right;
            while (cur.left != null) {
                cur = cur.left;
            }
            return cur;
        }

        /**
         * 查找节点
         *
         * @param key
         * @return
         */
        private RBNode<T> find(int key) {
            if (root == null) {
                return null;
            }
            RBNode<T> cur = root;
            while (cur != null) {
                if (key == cur.key) {
                    return cur;
                } else if (key < cur.key) {
                    cur = cur.left;
                } else {
                    cur = cur.right;
                }
            }
            return null;
        }

        /**
         * 添加节点
         *
         * @param key
         * @param value
         */
        public void put(int key, T value) {
            if (root == null) {
                root = new RBNode<>(key, value);
                root.color = Color.BLACK;
                return;
            }
            // 添加节点
            RBNode<T> cur = root;
            RBNode<T> parent = null;
            while (cur != null) {
                parent = cur;
                if (key == cur.key) {
                    cur.value = value;
                    return;
                }
                if (key < cur.key) {
                    cur = cur.left;
                } else {
                    cur = cur.right;
                }
            }
            cur = new RBNode<>(key, value);
            if (parent.key < key) {
                parent.right = cur;
            } else {
                parent.left = cur;
            }
            cur.parent = parent;
            // 调整红黑树
            adjustTree(cur);


        }

        /**
         * 调整红黑树
         *
         * @param node
         */
        private void adjustTree(RBNode<T> node) {
            if (node == root) { // 根节点为黑色，不需要调整
                node.color = Color.BLACK;
                return;
            }

            RBNode<T> parent = node.parent;
            RBNode<T> grandParent = parent.parent;
            RBNode<T> uncle = node.uncle();

            if (isBlack(parent)) { // 父节点为黑色，不需要调整
                return;
            }


            if (isRed(parent)) { // 父节点为红色，需要调整
                // 分两种情况
                // 1. 父节点为红色，叔叔节点为红色
                if (isRed(uncle)) {
                    // 将叔叔节点和父节点都置为黑色，祖父节点置为红色
                    parent.color = Color.BLACK;
                    uncle.color = Color.BLACK;
                    grandParent.color = Color.RED;
                    adjustTree(grandParent); // 递归调整祖父节点
                    return;
                }

                // 2. 父节点为红色，叔叔节点为黑色
                // 分4种情况
                // 2.1 父节点为左孩子，插入节点为父节点的左孩子
                if (parent.isLeft() && node.isLeft()) {
                    parent.color = Color.BLACK;
                    grandParent.color = Color.RED;
                    rightRotate(grandParent); // 右旋转,祖父节点
                    return;
                }
                // 2.2 父节点为左孩子，插入节点为父节点的右孩子
                if (parent.isLeft() && node.isRight()) {
                    leftRotate(parent); // 左旋转，父节点
                    node.color = Color.BLACK;
                    grandParent.color = Color.RED;
                    rightRotate(grandParent); // 右旋转，祖父节点
                    return;
                }
                // 2.3 父节点为右孩子，插入节点为父节点的左孩子
                if (parent.isRight() && node.isLeft()) {
                    rightRotate(parent); // 右旋转，父节点
                    node.color = Color.BLACK;
                    grandParent.color = Color.RED;
                    leftRotate(grandParent); // 左旋转，祖父节点
                    return;
                }
                // 2.4 父节点为右孩子，插入节点为父节点的右孩子
                if (parent.isRight() && node.isRight()) {
                    parent.color = Color.BLACK;
                    grandParent.color = Color.RED;
                    leftRotate(grandParent); // 左旋转，祖父节点
                    return;
                }

            }
        }

        /**
         * 左旋转
         *
         * @param node
         */
        private void leftRotate(RBNode<T> node) {
            RBNode<T> parent = node.parent;
            RBNode<T> right = node.right;
            RBNode<T> rl = right.left;
            if (node.isLeft()) {
                parent.left = right;
            } else {
                parent.right = right;
            }
            right.parent = parent;
            right.left = node;
            node.right = rl;
            node.parent = right;
        }

        /**
         * 右旋转
         *
         * @param node
         */
        private void rightRotate(RBNode<T> node) {
            RBNode<T> parent = node.parent;
            RBNode<T> left = node.left;
            RBNode<T> lr = left.right;

            if (node.isLeft()) {
                parent.left = left;
            } else {
                parent.right = left;
            }
            left.right = node;
            left.parent = parent;
            node.left = lr;
            node.parent = left;
        }

        /**
         * 是否为红色
         *
         * @param node
         * @return
         */
        private boolean isRed(RBNode<T> node) {
            return node != null && node.color == Color.RED;
        }

        /**
         * 是否为黑色
         *
         * @param node
         * @return
         */
        private boolean isBlack(RBNode<T> node) {
            return node == null || node.color == Color.BLACK;
        }

        enum Color {
            RED(true), BLACK(false); // 红色为true，黑色为false
            private boolean isRed;

            Color(boolean isRed) {
                this.isRed = isRed;
            }

            Color() {

            }

            public boolean isRed() {
                return isRed;
            }
        }

        /**
         * 红黑树节点
         *
         * @param <T>
         */
        private static class RBNode<T> {
            private RBNode<T> left;
            private RBNode<T> right;
            private RBNode<T> parent;
            private int key;
            private T value;
            private Color color;

            public RBNode(int key, T value) {
                this.key = key;
                this.value = value;
                this.color = Color.RED;
            }

            public RBNode(RBNode<T> left, RBNode<T> right, RBNode<T> parent, int key, T value, Color color) {
                this.left = left;
                this.right = right;
                this.parent = parent;
                this.key = key;
                this.value = value;
                this.color = color;
            }

            public boolean isLeft() {
                return parent != null && this == parent.left;
            }

            public boolean isRight() {
                return parent != null && this == parent.right;
            }

            /**
             * 获取兄弟节点
             *
             * @return
             */
            public RBNode<T> sibling() {
                if (isLeft()) {
                    return parent.right;
                }
                if (isRight()) {
                    return parent.left;
                }
                return null;
            }

            /**
             * 获取叔叔节点
             *
             * @return
             */
            public RBNode<T> uncle() {
                if (parent == null || parent.parent == null) {
                    return null;
                }
                if (parent.isLeft()) {
                    return parent.parent.right;
                }
                return parent.parent.left;
            }

        }
    }

}
